# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.

# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
# --------------------------------------------------------
# References:
# DeiT: https://github.com/facebookresearch/deit
# --------------------------------------------------------

import os
import PIL

from torchvision import datasets, transforms

from timm.data import create_transform
from timm.data.constants import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD



# 这是源代码函数
# def build_dataset(is_train, args):
#     transform = build_transform(is_train, args)
#
#     root = os.path.join(args.data_path, 'train' if is_train else 'val')
#     dataset = datasets.ImageFolder(root, transform=transform)
#     # dataset = datasets.ImageFolder(root)
#
#     print(dataset)
#
#     return dataset



def build_dataset(is_train=True, args=None):
    """
    构建数据集，支持训练和验证数据集的创建。

    参数：
        is_train (bool): 是否为训练集，默认为 True。
        args (Namespace): 参数对象，包含数据路径、输入大小等配置信息。

    返回：
        dataset (torchvision.datasets.ImageFolder): 返回构建好的数据集。
    """
    if is_train:
        # 训练集的图像预处理流程：
        # 这里注释掉了原本的随机裁剪和水平翻转操作，改为固定的 resize 和 center crop
        transform = transforms.Compose([
            transforms.Resize(args.input_size),  # 将图像调整为指定的输入大小
            transforms.CenterCrop(args.input_size),  # 从中心进行裁剪
            transforms.ToTensor(),  # 将图像转换为 Tensor 格式
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化处理
        ])
    else:
        # 验证集的图像预处理流程：
        transform = transforms.Compose([
            transforms.Resize(args.input_size),  # 将图像调整为指定的输入大小
            transforms.CenterCrop(args.input_size),  # 从中心进行裁剪
            transforms.ToTensor(),  # 将图像转换为 Tensor 格式
            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 归一化处理
        ])

    # 根据 is_train 标志选择训练集或验证集的路径
    root = os.path.join(args.data_path, 'train' if is_train else 'val')
    # 使用 ImageFolder 构建数据集，ImageFolder 会根据文件夹名称自动标注类别
    dataset = datasets.ImageFolder(root, transform=transform)
    return dataset


def build_transform(is_train, args):
    """
    构建训练或验证数据的图像预处理流水线。

    参数：
        is_train (bool): 是否为训练集，默认为 True。
        args (Namespace): 参数对象，包含数据路径、输入大小等配置信息。

    返回：
        transform (torchvision.transforms.Compose): 返回预处理流程。
    """
    # 定义 ImageNet 的均值和标准差（可根据需求调整）
    mean = (0.160377, 0.160377, 0.160377)
    std = (0.13723223, 0.13723223, 0.13723223)

    # 如果是训练集
    if is_train:
        # 构建训练集的变换（包括数据增强等）
        transform = create_transform(
            input_size=args.input_size,  # 输入图像的大小
            is_training=True,  # 是否为训练模式
            color_jitter=args.color_jitter,  # 颜色抖动
            auto_augment=args.aa,  # 自动增强策略
            interpolation='bicubic',  # 图像插值方法，使用双三次插值
            re_prob=args.reprob,  # 重标定概率
            re_mode=args.remode,  # 重标定模式
            re_count=args.recount,  # 重标定次数
            mean=mean,  # 均值
            std=std,  # 标准差
        )
        # 返回训练集的预处理流水线
        return transform

    # 如果是验证集
    t = []
    # 根据输入大小和裁剪比例来调整图像大小
    if args.input_size <= 224:
        crop_pct = 224 / 256  # 如果输入大小小于等于224，设置一个裁剪比例
    else:
        crop_pct = 1.0  # 如果输入大小大于224，裁剪比例为1.0
    size = int(args.input_size / crop_pct)  # 计算目标尺寸
    t.append(
        transforms.Resize(size, interpolation=PIL.Image.BICUBIC),  # 调整图像大小，保持与224图像相同的宽高比
    )
    t.append(transforms.CenterCrop(args.input_size))  # 从中心裁剪成指定的输入大小

    t.append(transforms.ToTensor())  # 将图像转换为 Tensor
    t.append(transforms.Normalize(mean, std))  # 归一化图像数据
    return transforms.Compose(t)  # 返回最终的图像预处理流水线
